package com.niodata.dp.monitor.components.task;

import com.niodata.dp.core.DpConfiguration;
import com.niodata.dp.core.RedisConstants;
import com.niodata.dp.log.LogUtil;
import com.niodata.dp.monitor.common.configure.MonitorConfigurationKeys;
import com.niodata.dp.monitor.common.service.RedisPool;
import com.niodata.dp.monitor.core.task.dao.JobLogMappingMonitorMapper;
import com.niodata.dp.monitor.core.task.entity.JobLogMapping;
import com.niodata.dp.monitor.core.task.entity.LogType;
import com.niodata.dp.util.RedisUtil;
import jakarta.annotation.PostConstruct;
import java.io.File;
import java.io.IOException;
import java.sql.Timestamp;
import java.util.Arrays;
import java.util.Date;
import java.util.EnumSet;
import java.util.List;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
import org.apache.hadoop.yarn.api.records.ApplicationAttemptReport;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.ApplicationReport;
import org.apache.hadoop.yarn.api.records.YarnApplicationState;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.log4j.Logger;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import redis.clients.jedis.JedisPool;

@Service("jobLogManagerImpl")
@Transactional
public class YarnJobLogManagerImpl {

  private static final Logger logger = Logger.getLogger(YarnJobLogManagerImpl.class);

  @Value("${" + MonitorConfigurationKeys.HADOOP_USER + "}")
  private String loginUser;

  @Value("${" + MonitorConfigurationKeys.HADOOP_USER_KEYTAB + ":NONE}")
  private String keyTabPath;

  @Value("${" + MonitorConfigurationKeys.HADOOP_HOME + "}")
  private String hadoopHome;

  @Autowired
  private JobLogMappingMonitorMapper jobLogMappingMapper;

  private YarnClient yarnClient;

  private JedisPool jedisPool = RedisPool.getInstance().getJedisPool();

  private static ApplicationId parseApplicationId(String applicationId) {
    //application_1558060967614_4609
    String[] ss = applicationId.split("_");
    return ApplicationId.newInstance(Long.parseLong(ss[1]), Integer.parseInt(ss[2]));
  }

  @PostConstruct
  public void init() {
    if (DpConfiguration.isHadoopEnabled()) {
      this.yarnClient = getYarnClient();
      new Thread(new YarnLogConsumeThread()).start();
      new Thread(new YarnAppStatusCheckThread()).start();
    }
  }

  private YarnClient getYarnClient() {
    YarnClient client = YarnClient.createYarnClient();
    Configuration configuration = new Configuration();
    File[] configFiles = new File(hadoopHome + "/etc/hadoop").listFiles();
    Arrays.stream(configFiles).forEach(configFile -> {
      if (configFile.getAbsolutePath().endsWith(".xml")) {
        configuration.addResource(new Path(configFile.getAbsolutePath()));
      }
    });
    UserGroupInformation.setConfiguration(configuration);
    try {
      if (DpConfiguration.isHadoopKerberosEnabled()) {
        UserGroupInformation.loginUserFromKeytab(loginUser, keyTabPath);
      } else {
        UserGroupInformation.setLoginUser(UserGroupInformation.createRemoteUser(loginUser));
      }
    } catch (IOException e) {
      LogUtil.logStackTrace(logger, e);
    }
    client.init(configuration);
    client.start();
    return client;
  }

  private JobLogMapping parseLogInfo(String info) {
    String[] infos = info.split(";");
    JobLogMapping jobLogMapping = new JobLogMapping();
    for (String item : infos) {
      String[] values = item.split("=");
      switch (values[0]) {
        case "logDir":
          jobLogMapping.setLogDir(values[1]);
          break;
        case "appId":
          jobLogMapping.setYarnAppId(values[1]);
          break;
        case "execId":
          jobLogMapping.setExecId(values[1]);
          break;
        case "taskName":
          jobLogMapping.setTaskName(values[1]);
          break;
        case "jobName":
          jobLogMapping.setJobName(values[1]);
          break;
        case "hostName":
          jobLogMapping.setHostName(values[1]);
          break;
        case "hostIp":
          jobLogMapping.setHostAddr(values[1]);
          break;
        default:
      }
    }
    jobLogMapping.setLogType(LogType.YARNLOG);
    jobLogMapping.setCreateTime(new Timestamp(new Date().getTime()));
    return jobLogMapping;
  }

  private ApplicationReport getApplicationReport(String appId) throws Exception {

    String[] appidArray = appId.split("_");
    if (appidArray == null || appidArray.length != 3) {
      logger.error("appid length error," + appId);
      return null;
    }
    ApplicationId applicationId = ApplicationId
          .newInstance(Long.valueOf(appidArray[1]), Integer.valueOf(appidArray[2]));
    if (null != applicationId) {
      return yarnClient.getApplicationReport(applicationId);
    }
    return null;
  }

  /**
   * sync yarn application log from redis to db.
   */
  class YarnLogConsumeThread implements Runnable {

    @Override
    public void run() {
      while (true) {
        try {
          String logInfo = RedisUtil.rpop(jedisPool,
                RedisConstants.DB_INDEX_YARN_LOG, RedisConstants.KEY_YARN_LOG_MAPPING);
          if (logInfo == null) {
            try {
              Thread.sleep(1500);
            } catch (InterruptedException e) {
              e.printStackTrace();
              return;
            }
            continue;
          }
          JobLogMapping jobLogMapping = parseLogInfo(logInfo);
          logger
                .info("fetch yarn app:" + jobLogMapping.getYarnAppId() + ",hostName:"
                      + jobLogMapping
                      .getHostName());
          ApplicationId applicationId = parseApplicationId(jobLogMapping.getYarnAppId());
          ApplicationAttemptId applicationAttemptId = ApplicationAttemptId
                .newInstance(applicationId, 1);
          jobLogMapping.setYarnAttemptId(applicationAttemptId.toString());
          jobLogMappingMapper.insert(jobLogMapping);
        } catch (Exception e) {
          e.printStackTrace();
        }
      }
    }

  }

  /**
   * check application attempts,some attempts my be lost because of polling interval.
   */
  class YarnAppStatusCheckThread implements Runnable {

    @Override
    public void run() {
      EnumSet<YarnApplicationState> applicationStates =
            EnumSet.of(YarnApplicationState.RUNNING);
      while (true) {
        try {
          List<ApplicationReport> reports = yarnClient.getApplications(applicationStates);
          for (ApplicationReport report : reports) {
            ApplicationAttemptId applicationAttemptId = report.getCurrentApplicationAttemptId();
            if (applicationAttemptId.getAttemptId() == 1) {
              continue;
            }
            String appId = report.getApplicationId().toString();
            JobLogMapping lastLog = jobLogMappingMapper.getYarnLatestJobLogByAppId(appId);
            if (lastLog == null) {
              logger.error("yarn log not found for application in db:" + appId);
              continue;
            }
            ApplicationAttemptReport attemptReport =
                  yarnClient.getApplicationAttemptReport(applicationAttemptId);
            if (lastLog.getYarnAttemptId().equals(applicationAttemptId.toString())) {
              continue;
            }
            lastLog.setHostName(attemptReport.getHost());
            lastLog.setHostAddr(attemptReport.getHost());
            lastLog.setYarnAttemptId(applicationAttemptId.toString());
            lastLog.setCreateTime(new Timestamp(System.currentTimeMillis()));
            lastLog.setYarnAppId(appId);
            String containerId = attemptReport.getAMContainerId().toString();
            String lastLogPath = lastLog.getLogDir();
            String logPath =
                  lastLogPath.substring(0, lastLogPath.lastIndexOf("/") + 1) + containerId;
            lastLog.setLogDir(logPath);
            lastLog.setId(0);
            jobLogMappingMapper.insert(lastLog);
          }
          Thread.sleep(20000);
        } catch (InterruptedException e) {
          return;
        } catch (Exception e) {
          try {
            Thread.sleep(20000);
          } catch (Exception e2) {
            //nothing to do
          }
        }
      }
    }
  }
}


